using namespace std;
#include<iostream>
#include<fstream>
#include<vector>
#include<stdio.h>
#include<stdlib.h>
#include<map>
#include<set>
#include"BA2.h"

extern std::vector<string> cmdList;
extern std::vector<Indiv> sample;
extern std::map<string,int> popToindex;
extern std::map<int,string> indexTopop;
extern std::map<string,int> labelToindex;
extern std::map<int,string> indexTolabel;
extern std::set<string> labelSet;
extern std::set<string> popSet;
extern std::vector< vector< vector<long int> > > alleleCounts;
extern std::vector< set<string> > alleleSet;
extern std::vector< map<string,int> > alleleMap;
extern std::vector< map<int,string> > intToallele;

// constructor for Parameters class

Parameters::Parameters()
{
  outfile="BA2.out";
  MCMCreps=10000;
  MCMCthin=100;
}

int TranslateCommand(string cmd)
{
  // parse command and options into vector of strings
  cmdList.erase(cmdList.begin(),cmdList.end());
  int pos=0; int prev_pos=0;
  while((pos=cmd.find_first_of(' ',pos))!=string::npos)
    {
      cmdList.push_back(cmd.substr(prev_pos,pos-prev_pos));
      prev_pos = ++pos;
    }
  cmdList.push_back(cmd.substr(prev_pos,pos-prev_pos));

  // translate command into integer
  if(cmdList[0]=="quit") return(1);
  else
    if(cmdList[0]=="scan") return(2);
    else
      if(cmdList[0]=="pset") return(3);
      else
	if(cmdList[0]=="help") return(4);
	else
	  if(cmdList[0]=="printData") return(5);
	  else
	    if(cmdList[0]=="printAlleles") return(6);
	    else
	      return(0);
}

void PrintBanner()
{
  cout << "             BayesAss V.2.0 [FlatHead Lake Edition]              " << "\n";
  cout << "               Bruce Rannala (http://rannala.org)                " << "\n";     
  cout << "         Genome Center, University of California, Davis          " <<  "\n";
  cout << "            Copyright 2007, University of California             " << "\n";
  cout << "               [enter \"help;\" to list commands]                " << "\n";
}

void PrintCommands()
{
	  cout << "                                                            \n";
	  cout << "------------------------------------------------------------\n";
	  cout << "               BayesAss V.2.0: List of Commands             \n";
	  cout << "------------------------------------------------------------\n";
	  cout << " citation [-C citation format] : articles to cite for BA2   \n";
	  cout << " help : print list of commands to screen                    \n";
	  cout << " printAlleles : print allele counts                         \n";
	  cout << " printData : print to screen data scanned from input file   \n";
	  cout << " pset [parameter=value] : set parameter to value            \n";
	  cout << " scan [-F file format][input file name] : open input file   \n";
	  cout << " quit : terminate the program                               \n";
	  cout << "------------------------------------------------------------\n";
	  cout << " Important note: all commands must be terminated with \";\" \n";
	  cout << "------------------------------------------------------------\n";
	  cout << "                                                            \n";
}

void DataSummary(Parameters pvalues)
{
  // print summary to screen
  cout << "Data File in " << pvalues.fileFormat << " format \n";
  cout << "Individuals: " << sample.size() << " ";
  cout << "Populations: " << popSet.size() << " ";
  cout << "Loci: " << sample[0].genotypes.size() << "\n";
}

void ReadData(ifstream &infile, Parameters &pvalues)
{
  string aLine;
  Gene geno;
  Indiv temp;

  cout << "Scanning input file: " << cmdList[2] << "\n";
  sample.erase(sample.begin(),sample.end());
  popSet.erase(popSet.begin(),popSet.end());
  popToindex.erase(popToindex.begin(),popToindex.end());
  indexTopop.erase(indexTopop.begin(),indexTopop.end());
  labelToindex.erase(labelToindex.begin(),labelToindex.end());
  indexTolabel.erase(indexTolabel.begin(),indexTolabel.end());
  alleleCounts.erase(alleleCounts.begin(),alleleCounts.end());
  alleleSet.erase(alleleSet.begin(),alleleSet.end());
  alleleMap.erase(alleleMap.begin(),alleleMap.end());
  intToallele.erase(intToallele.begin(),intToallele.end());
  labelSet.erase(labelSet.begin(),labelSet.end());
  pvalues.alist.erase(pvalues.alist.begin(),pvalues.alist.end());

  int i=0, popindex=0, indivindex=0;

  if(pvalues.fileFormat=="-FS")      // data is in Structure file format
    {
      while(getline(infile,aLine))
	{
	  i++;
	  vector<string> parsedLine;
	  int pos=0; int prev_pos=0;
	  while((pos=aLine.find_first_of(' ',pos))!=string::npos)
	    {
	      parsedLine.push_back(aLine.substr(prev_pos,pos-prev_pos));
	      prev_pos = ++pos;
	    }
	  parsedLine.push_back(aLine.substr(prev_pos,pos-prev_pos));
	  if((i % 2)!=0)
	    {
	      temp.label=parsedLine[0];
	      labelSet.insert(parsedLine[0]);
	      labelToindex[string(parsedLine[0])]=indivindex;
	      indexTolabel[indivindex]=parsedLine[0];
	      ++indivindex;
	      temp.sampledPopulation=parsedLine[1];
	      if(popSet.count(parsedLine[1])==0)
		{
		  popSet.insert(parsedLine[1]);
		  popToindex[string(parsedLine[1])]=popindex;
		  indexTopop[popindex]=parsedLine[1];
		  ++popindex;
		}
	      for(int j=2;j<=parsedLine.size()-1;j++)
		{ 
		  geno.g1=parsedLine[j];
		  temp.genotypes.push_back(geno);
		} 
	    }
	  else
	    {
	      for(int j=2;j<=parsedLine.size()-1;j++)
		{ 
		  temp.genotypes[j-2].g2=parsedLine[j];
		} 
	      sample.push_back(temp);
	      temp.genotypes.erase(temp.genotypes.begin(),temp.genotypes.end());
	    }
	} 
    } 
  GetAlleleCounts();
  DataSummary(pvalues);
}

void PrintData()
{
  for(int i=0;i<=sample.size()-1;i++)
    {
      cout << sample[i].label << " " << sample[i].sampledPopulation << " ";
      for(int j=0;j<=sample[i].genotypes.size()-1;j++)
	cout << sample[i].genotypes[j].g1 << "/" << sample[i].genotypes[j].g2 << " ";
      cout << "\n";
    }
}

void PrintParameters(Parameters pvalues)
{
	  cout << "                                                            \n";
	  cout << "------------------------------------------------------------\n";
	  cout << "               BayesAss V.2.0: Parameter Settings           \n";
	  cout << "------------------------------------------------------------\n";
	  cout << " alist=(";
	  if(pvalues.alist.size()==0)
	    cout << ")\n";
	  else
	    for(int i=0;i<=pvalues.alist.size()-1;i++)
	      {
		cout << pvalues.alist[i];
		if(i!=(pvalues.alist.size()-1)) cout << ",";
		else cout << ")\n";
	      }
	  cout << " MCMCreps=" << pvalues.MCMCreps << "\n";
	  cout << " MCMCthin=" << pvalues.MCMCthin << "\n";
	  cout << " outfile=" << pvalues.outfile << "\n";
	  cout << "------------------------------------------------------------\n";
	  cout << "                                                            \n";
}

int SetParameters(Parameters &pvalues)
{
  string parm;
  string paramValue;
  int outfile_changed=0;
  if(cmdList.size()==1)
    {
      PrintParameters(pvalues);
    }
  else
    if(cmdList.size()==2)
      {
	int pos=cmdList[1].find_first_of("=");
	if(pos!=string::npos)
	  {
	    parm=cmdList[1].substr(0,pos);
	    paramValue=cmdList[1].substr(pos+1);
	  }
	if(parm=="MCMCreps") pvalues.MCMCreps=atol(paramValue.c_str());
	else
	  if(parm=="MCMCthin") pvalues.MCMCthin=atol(paramValue.c_str());
	  else
	    if(parm=="outfile") {pvalues.outfile=paramValue; outfile_changed=1;}
	    else
	      if(parm=="alist") 
		{ 
		  vector<string> tempList;
		  paramValue.erase(paramValue.begin(),paramValue.begin()+1);
		  int pos=0; int prev_pos=0;
		  while((pos=paramValue.find_first_of(',',pos))!=string::npos)
		    {
		      tempList.push_back(paramValue.substr(prev_pos,pos-prev_pos));
		      prev_pos = ++pos;
		    }
		  tempList.push_back(paramValue.substr(prev_pos,pos-prev_pos));
		  tempList[tempList.size()-1]= 
		    tempList[tempList.size()-1].erase(tempList[tempList.size()-1].size()-1,tempList[tempList.size()-1].size());

		  // check that individuals exist in sample, otherwise return error
		  bool missingLabel=0;
		  for(int i=0;i<=tempList.size()-1;i++)
		    if(labelSet.count(tempList[i])==0) {cout << "Error: label " << tempList[i] << " not in sample\n"; missingLabel=1;} 
		  if(!missingLabel)
		    pvalues.alist=tempList;
		}
	PrintParameters(pvalues);
	return(outfile_changed);
      }
}

void GetAlleleCounts()
{

  // create map from allele labels to index numbers

  for(int i=0;i<=sample[0].genotypes.size()-1;i++)
    {
      set<string> tempSet;
      map<string,int> tempMap;
      map<int,string> tempintToallele;
      int alleleit=0;
      for(int j=0;j<=sample.size()-1;j++)
	{
	  if(tempSet.count(sample[j].genotypes[i].g1)==0)
	    {
	      tempSet.insert(sample[j].genotypes[i].g1);
	      tempMap[string(sample[j].genotypes[i].g1)]=alleleit;
	      tempintToallele[alleleit]=sample[j].genotypes[i].g1;
	      ++alleleit;
	    }
	  if(tempSet.count(sample[j].genotypes[i].g2)==0)
	    {
	      tempSet.insert(sample[j].genotypes[i].g2);
	      tempMap[string(sample[j].genotypes[i].g2)]=alleleit;
	      tempintToallele[alleleit]=sample[j].genotypes[i].g2;
	      ++alleleit;
	    }
	}
      alleleSet.push_back(tempSet);
      alleleMap.push_back(tempMap);
      intToallele.push_back(tempintToallele);
    }

  // create data structure for allele counts

  vector< vector<long int> > tempPopln;
  for(int j=0;j<=sample[0].genotypes.size()-1;j++)
    {
      vector<long int> tempLocus;
      for(int k=0; k<=alleleSet[j].size()-1; k++)
	tempLocus.push_back(0);
      tempPopln.push_back(tempLocus);
    }
  for(int i=0;i<=popSet.size()-1;i++)
    {
      alleleCounts.push_back(tempPopln);
    }

  // tabulate allele counts
  for(int i=0;i<=sample.size()-1;i++)
    for(int j=0;j<=sample[0].genotypes.size()-1;j++)
      {
	alleleCounts[popToindex[string(sample[i].sampledPopulation)]][j][alleleMap[j][string(sample[i].genotypes[j].g1)]]++;
	alleleCounts[popToindex[string(sample[i].sampledPopulation)]][j][alleleMap[j][string(sample[i].genotypes[j].g2)]]++;
      }
}

void PrintAlleleCounts()
{
  for(int i=0;i<=popSet.size()-1;i++)
    {
      cout << "\nPopulation: " << indexTopop[i] << "\n";
      for(int j=0;j<=alleleCounts[i].size()-1;j++)
	{
	  cout << "Locus " << j << ": ";
	  for(int k=0;k<=alleleCounts[i][j].size()-1;k++)
	    cout << intToallele[j][k]<< ":" << alleleCounts[i][j][k] << " ";
	  cout << "\n";
	}
    } 
  cout << "\n";
}
